{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "1. Loads vocab\n",
    "2. Builds data generator\n",
    "3. Builds image caption model\n",
    "4. Train\n",
    "4. Eval\n",
    "\"\"\"\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import tensorflow as tf\n",
    "from tensorflow import gfile\n",
    "from tensorflow import logging\n",
    "import pprint\n",
    "import cPickle\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "input_description_file = \"../../image_caption_data/results_20130124.token\"\n",
    "input_img_feature_dir = \"../../image_caption_data/feature_extraction_inception_v3\"\n",
    "input_vocab_file = \"../../image_caption_data/vocab.txt\"\n",
    "output_dir = \"../../image_caption_data/local_run\"\n",
    "\n",
    "if not gfile.Exists(output_dir):\n",
    "    gfile.MakeDirs(output_dir)\n",
    "\n",
    "def get_default_params():\n",
    "    return tf.contrib.training.HParams(\n",
    "        num_vocab_word_threshold=3,\n",
    "        num_embedding_nodes=32,\n",
    "        num_timesteps=10,\n",
    "        num_lstm_nodes=[64, 64],\n",
    "        num_lstm_layers=2,\n",
    "        num_fc_nodes=32,\n",
    "        batch_size=50,\n",
    "        cell_type='lstm',\n",
    "        clip_lstm_grads=1.0,\n",
    "        learning_rate=0.001,\n",
    "        keep_prob=0.8,\n",
    "        log_frequent=100,\n",
    "        save_frequent=1000,\n",
    "    )\n",
    "\n",
    "hps = get_default_params()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Vocab(object):\n",
    "    def __init__(self, filename, word_num_threshold):\n",
    "        self._id_to_word = {}\n",
    "        self._word_to_id = {}\n",
    "        self._unk = -1\n",
    "        self._eos = -1\n",
    "        self._word_num_threshold = word_num_threshold\n",
    "        self._read_dict(filename)\n",
    "\n",
    "    def _read_dict(self, filename):\n",
    "        with gfile.GFile(filename, 'r') as f:\n",
    "            lines = f.readlines()\n",
    "        for line in lines:\n",
    "            word, occurence = line.strip('\\r\\n').split('\\t')\n",
    "            occurence = int(occurence)\n",
    "            if word != '<UNK>' and occurence < self._word_num_threshold:\n",
    "                continue\n",
    "            idx = len(self._id_to_word)\n",
    "            if word == '<UNK>':\n",
    "                self._unk = idx\n",
    "            elif word == '.':\n",
    "                self._eos = idx\n",
    "            if idx in self._id_to_word or word in self._word_to_id:\n",
    "                raise Exception('duplicate words in vocab file')\n",
    "            self._word_to_id[word] = idx\n",
    "            self._id_to_word[idx] = word\n",
    "\n",
    "    @property\n",
    "    def unk(self):\n",
    "        return self._unk\n",
    "\n",
    "    @property\n",
    "    def eos(self):\n",
    "        return self._eos\n",
    "\n",
    "    def word_to_id(self, word):\n",
    "        return self._word_to_id.get(word, self.unk)\n",
    "\n",
    "    def id_to_word(self, cur_id):\n",
    "        return self._id_to_word.get(cur_id, '<UNK>')\n",
    "\n",
    "    def size(self):\n",
    "        return len(self._word_to_id)\n",
    "\n",
    "    def encode(self, sentence):\n",
    "        word_ids = [self.word_to_id(cur_word) for cur_word in sentence.split(' ')]\n",
    "        return word_ids\n",
    "\n",
    "    def decode(self, sentence_id):\n",
    "        words = [self.id_to_word(word_id) for word_id in sentence_id]\n",
    "        return ' '.join(words)\n",
    "    \n",
    "\n",
    "def parse_token_file(token_file):\n",
    "    \"\"\"Parses token file.\"\"\"\n",
    "    img_name_to_tokens = {}\n",
    "    with gfile.GFile(token_file, 'r') as f:\n",
    "        lines = f.readlines()\n",
    "    for line in lines:\n",
    "        img_id, description = line.strip('\\r\\n').split('\\t')\n",
    "        img_name, _ = img_id.split('#')\n",
    "        img_name_to_tokens.setdefault(img_name, [])\n",
    "        img_name_to_tokens[img_name].append(description)\n",
    "    return img_name_to_tokens\n",
    "\n",
    "def convert_token_to_id(img_name_to_tokens, vocab):\n",
    "    \"\"\"Converts tokens of each description of imgs to id. \"\"\"\n",
    "    img_name_to_token_ids = {}\n",
    "    for img_name in img_name_to_tokens:\n",
    "        img_name_to_token_ids.setdefault(img_name, [])\n",
    "        descriptions = img_name_to_tokens[img_name]\n",
    "        for description in descriptions:\n",
    "            token_ids = vocab.encode(description)\n",
    "            img_name_to_token_ids[img_name].append(token_ids)\n",
    "    return img_name_to_token_ids\n",
    "\n",
    "vocab = Vocab(input_vocab_file, hps.num_vocab_word_threshold)\n",
    "vocab_size = vocab.size()\n",
    "logging.info(\"vocab_size: %d\" % vocab_size)\n",
    "    \n",
    "    \n",
    "img_name_to_tokens = parse_token_file(input_description_file)\n",
    "img_name_to_token_ids = convert_token_to_id(img_name_to_tokens, vocab)\n",
    "\n",
    "logging.info(\"num of all images: %d\" % len(img_name_to_tokens))\n",
    "pprint.pprint(img_name_to_tokens.keys()[0:10])\n",
    "pprint.pprint(img_name_to_tokens['2778832101.jpg'])\n",
    "logging.info(\"num of all images: %d\" % len(img_name_to_token_ids))\n",
    "pprint.pprint(img_name_to_token_ids.keys()[0:10])\n",
    "pprint.pprint(img_name_to_token_ids['2778832101.jpg'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImageCaptionData(object):\n",
    "    def __init__(self,\n",
    "                 img_name_to_token_ids,\n",
    "                 img_feature_dir,\n",
    "                 num_timesteps,\n",
    "                 vocab,\n",
    "                 deterministic = False):\n",
    "        self._vocab = vocab\n",
    "        self._all_img_feature_filepaths = []\n",
    "        for filename in gfile.ListDirectory(img_feature_dir):\n",
    "            self._all_img_feature_filepaths.append(os.path.join(img_feature_dir, filename))\n",
    "        pprint.pprint(self._all_img_feature_filepaths)\n",
    "\n",
    "        self._img_name_to_token_ids = img_name_to_token_ids\n",
    "        self._num_timesteps = num_timesteps\n",
    "        self._indicator = 0\n",
    "        self._deterministic = deterministic\n",
    "        self._img_feature_filenames = []\n",
    "        self._img_feature_data = []\n",
    "        self._load_img_feature_pickle()\n",
    "        if not self._deterministic:\n",
    "            self._random_shuffle()\n",
    "\n",
    "\n",
    "    def _load_img_feature_pickle(self):\n",
    "        for filepath in self._all_img_feature_filepaths:\n",
    "            logging.info(\"loading %s\" % filepath)\n",
    "            with gfile.GFile(filepath, 'r') as f:\n",
    "                filenames, features = cPickle.load(f)\n",
    "                self._img_feature_filenames += filenames\n",
    "                self._img_feature_data.append(features)\n",
    "        self._img_feature_data = np.vstack(self._img_feature_data)\n",
    "        origin_shape = self._img_feature_data.shape\n",
    "        self._img_feature_data = np.reshape(\n",
    "            self._img_feature_data, (origin_shape[0], origin_shape[3]))\n",
    "        self._img_feature_filenames = np.asarray(self._img_feature_filenames)\n",
    "        print(self._img_feature_data.shape)\n",
    "        print(self._img_feature_filenames.shape)\n",
    "        if not self._deterministic:\n",
    "            self._random_shuffle()\n",
    "\n",
    "\n",
    "    def size(self):\n",
    "        return len(self._img_feature_filenames)\n",
    "\n",
    "    def img_feature_size(self):\n",
    "        return self._img_feature_data.shape[1]\n",
    "\n",
    "    def _random_shuffle(self):\n",
    "        p = np.random.permutation(self.size())\n",
    "        self._img_feature_filenames = self._img_feature_filenames[p]\n",
    "        self._img_feature_data = self._img_feature_data[p]\n",
    "\n",
    "    def _img_desc(self, filenames):\n",
    "        batch_sentence_ids = []\n",
    "        batch_weights = []\n",
    "        for filename in filenames:\n",
    "            token_ids_set = self._img_name_to_token_ids[filename]\n",
    "            # chosen_token_ids = random.choice(token_ids_set)\n",
    "            chosen_token_ids = token_ids_set[0]\n",
    "            chosen_token_length = len(chosen_token_ids)\n",
    "\n",
    "            weight = [1 for i in range(chosen_token_length)]\n",
    "            if chosen_token_length >= self._num_timesteps:\n",
    "                chosen_token_ids = chosen_token_ids[0:self._num_timesteps]\n",
    "                weight = weight[0:self._num_timesteps]\n",
    "            else:\n",
    "                remaining_length = self._num_timesteps - chosen_token_length\n",
    "                chosen_token_ids += [self._vocab.eos for i in range(remaining_length)]\n",
    "                weight += [0 for i in range(remaining_length)]\n",
    "            batch_sentence_ids.append(chosen_token_ids)\n",
    "            batch_weights.append(weight)\n",
    "        batch_sentence_ids = np.asarray(batch_sentence_ids)\n",
    "        batch_weights = np.asarray(batch_weights)\n",
    "        return batch_sentence_ids, batch_weights\n",
    "\n",
    "    def next(self, batch_size):\n",
    "        end_indicator = self._indicator + batch_size\n",
    "        if end_indicator > self.size():\n",
    "            if not self._deterministic:\n",
    "                self._random_shuffle()\n",
    "            self._indicator = 0\n",
    "            end_indicator = self._indicator + batch_size\n",
    "        assert end_indicator <= self.size()\n",
    "\n",
    "        batch_img_features = self._img_feature_data[self._indicator: end_indicator]\n",
    "        batch_img_names = self._img_feature_filenames[self._indicator: end_indicator]\n",
    "        batch_sentence_ids, batch_weights = self._img_desc(batch_img_names)\n",
    "\n",
    "        self._indicator = end_indicator\n",
    "        return batch_img_features, batch_sentence_ids, batch_weights, batch_img_names\n",
    "\n",
    "\n",
    "caption_data = ImageCaptionData(img_name_to_token_ids, input_img_feature_dir, hps.num_timesteps, vocab)\n",
    "img_feature_dim = caption_data.img_feature_size()\n",
    "caption_data_size = caption_data.size()\n",
    "logging.info(\"img_feature_dim: %d\" % img_feature_dim)\n",
    "logging.info(\"caption_data_size: %d\" % caption_data_size)\n",
    "\n",
    "batch_img_features, batch_sentence_ids, batch_weights, batch_img_names = caption_data.next(5)\n",
    "pprint.pprint(batch_img_features)\n",
    "pprint.pprint(batch_sentence_ids)\n",
    "pprint.pprint(batch_weights)\n",
    "pprint.pprint(batch_img_names)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_rnn_cell(hidden_dim, cell_type):\n",
    "    if cell_type == 'lstm':\n",
    "        return tf.contrib.rnn.BasicLSTMCell(hidden_dim, state_is_tuple=True)\n",
    "    elif cell_type == 'gru':\n",
    "        return tf.contrib.rnn.GRUCell(hidden_dim)\n",
    "    else:\n",
    "        raise Exception(\"%s has not been supported\" % cell_type)\n",
    "\n",
    "def dropout(cell, keep_prob):\n",
    "    return tf.contrib.rnn.DropoutWrapper(cell, output_keep_prob=keep_prob)\n",
    "\n",
    "\n",
    "def get_train_model(hps, vocab_size, img_feature_dim):\n",
    "    num_timesteps = hps.num_timesteps\n",
    "    batch_size = hps.batch_size\n",
    "\n",
    "    img_feature  = tf.placeholder(tf.float32, (batch_size, img_feature_dim))\n",
    "    sentence = tf.placeholder(tf.int32, (batch_size, num_timesteps))\n",
    "    mask = tf.placeholder(tf.float32, (batch_size, num_timesteps))\n",
    "    keep_prob = tf.placeholder(tf.float32, name='keep_prob')\n",
    "\n",
    "    global_step = tf.Variable(tf.zeros([], tf.int64), name='global_step', trainable=False)\n",
    "\n",
    "    # Sets up the embedding layer.\n",
    "    embedding_initializer = tf.random_uniform_initializer(-1.0, 1.0)\n",
    "    with tf.variable_scope('embedding', initializer=embedding_initializer):\n",
    "        embeddings = tf.get_variable(\n",
    "            'embeddings',\n",
    "            [vocab_size, hps.num_embedding_nodes],\n",
    "            tf.float32)\n",
    "        embed_token_ids = tf.nn.embedding_lookup(embeddings, sentence[:, 0:num_timesteps-1])\n",
    "\n",
    "    img_feature_embed_init = tf.uniform_unit_scaling_initializer(factor=1.0)\n",
    "    with tf.variable_scope('image_feature_embed', initializer=img_feature_embed_init):\n",
    "        embed_img = tf.layers.dense(img_feature, hps.num_embedding_nodes)\n",
    "        embed_img = tf.expand_dims(embed_img, 1)\n",
    "        embed_inputs = tf.concat([embed_img, embed_token_ids], axis=1)\n",
    "\n",
    "    # Sets up LSTM network.\n",
    "    scale = 1.0 / math.sqrt(hps.num_embedding_nodes + hps.num_lstm_nodes[-1]) / 3.0\n",
    "    lstm_init = tf.random_uniform_initializer(-scale, scale)\n",
    "    with tf.variable_scope('lstm_nn', initializer=lstm_init):\n",
    "        cells = []\n",
    "        for i in range(hps.num_lstm_layers):\n",
    "            cell = create_rnn_cell(hps.num_lstm_nodes[i], hps.cell_type)\n",
    "            cell = dropout(cell, keep_prob)\n",
    "            cells.append(cell)\n",
    "        cell = tf.contrib.rnn.MultiRNNCell(cells)\n",
    "\n",
    "        initial_state = cell.zero_state(hps.batch_size, tf.float32)\n",
    "        # rnn_outputs: [batch_size, num_timesteps, hps.num_lstm_node[-1]]\n",
    "        rnn_outputs, _ = tf.nn.dynamic_rnn(cell,\n",
    "                                           embed_inputs,\n",
    "                                           initial_state=initial_state)\n",
    "\n",
    "    # Sets up the fully-connected layer.\n",
    "    fc_init = tf.uniform_unit_scaling_initializer(factor=1.0)\n",
    "    with tf.variable_scope('fc', initializer=fc_init):\n",
    "        rnn_outputs_2d = tf.reshape(rnn_outputs, [-1, hps.num_lstm_nodes[-1]])\n",
    "        fc1 = tf.layers.dense(rnn_outputs_2d, hps.num_fc_nodes, name='fc1')\n",
    "        fc1_dropout = tf.contrib.layers.dropout(fc1, keep_prob)\n",
    "        fc1_dropout = tf.nn.relu(fc1_dropout)\n",
    "        logits = tf.layers.dense(fc1_dropout, vocab_size, name='logits')\n",
    "\n",
    "    with tf.variable_scope('loss'):\n",
    "        sentence_flatten = tf.reshape(sentence, [-1])\n",
    "        mask_flatten = tf.reshape(mask, [-1])\n",
    "        mask_sum = tf.reduce_sum(mask_flatten)\n",
    "        softmax_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
    "            logits=logits, labels=sentence_flatten)\n",
    "        weighted_softmax_loss = tf.multiply(softmax_loss,\n",
    "                                            tf.cast(mask_flatten, tf.float32))\n",
    "        prediction = tf.argmax(logits, 1, output_type = tf.int32)\n",
    "        correct_prediction = tf.equal(prediction, sentence_flatten)\n",
    "        correct_prediction_with_mask = tf.multiply(\n",
    "            tf.cast(correct_prediction, tf.float32),\n",
    "            mask_flatten)\n",
    "        accuracy = tf.reduce_sum(correct_prediction_with_mask) / mask_sum\n",
    "        loss = tf.reduce_sum(weighted_softmax_loss) / mask_sum\n",
    "        tf.summary.scalar('loss', loss)\n",
    "\n",
    "    with tf.variable_scope('train_op'):\n",
    "        tvars = tf.trainable_variables()\n",
    "        for var in tvars:\n",
    "            logging.info(\"variable name: %s\" % (var.name))\n",
    "        grads, _ = tf.clip_by_global_norm(\n",
    "            tf.gradients(loss, tvars), hps.clip_lstm_grads)\n",
    "        for grad, var in zip(grads, tvars):\n",
    "            tf.summary.histogram('%s_grad' % (var.name), grad)\n",
    "        optimizer = tf.train.AdamOptimizer(hps.learning_rate)\n",
    "        train_op = optimizer.apply_gradients(zip(grads, tvars), global_step=global_step)\n",
    "\n",
    "    return ((img_feature, sentence, mask, keep_prob),\n",
    "            (loss, accuracy, train_op),\n",
    "            global_step)\n",
    "\n",
    "placeholders, metrics, global_step = get_train_model(hps, vocab_size, img_feature_dim)\n",
    "img_feature, sentence, mask, keep_prob = placeholders\n",
    "loss, accuracy, train_op = metrics\n",
    "\n",
    "summary_op = tf.summary.merge_all()\n",
    "\n",
    "init_op = tf.global_variables_initializer()\n",
    "saver = tf.train.Saver(max_to_keep=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_steps = 10000\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    sess.run(init_op)\n",
    "    writer = tf.summary.FileWriter(output_dir, sess.graph)\n",
    "    for i in range(training_steps):\n",
    "        batch_img_features, batch_sentence_ids, batch_weights, _ = caption_data.next(hps.batch_size)\n",
    "        input_vals = (batch_img_features, batch_sentence_ids, batch_weights, hps.keep_prob)\n",
    "        \n",
    "        feed_dict = dict(zip(placeholders, input_vals))\n",
    "        fetches = [global_step, loss, accuracy, train_op]\n",
    "        \n",
    "        should_log = (i + 1) % hps.log_frequent == 0\n",
    "        should_save = (i + 1) % hps.save_frequent == 0\n",
    "        if should_log:\n",
    "            fetches += [summary_op]\n",
    "        outputs = sess.run(fetches, feed_dict)\n",
    "        global_step_val, loss_val, accuracy_val = outputs[0:3]\n",
    "        if should_log:\n",
    "            summary_str = outputs[4]\n",
    "            writer.add_summary(summary_str, global_step_val)\n",
    "            logging.info('Step: %5d, loss: %3.3f, accuracy: %3.3f'\n",
    "                         % (global_step_val, loss_val, accuracy_val))\n",
    "        if should_save:\n",
    "            logging.info(\"Step: %d, image caption model saved\" % (global_step_val))\n",
    "            saver.save(sess, os.path.join(output_dir, \"image_caption\"), global_step=global_step_val)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
